#include <iostream>
#include <iomanip>
#include <vector>
#include <gmssl/pbkdf2.h>
#include <gmssl/rand.h>
#include <gmssl/sm9.h>
#include <gmssl/sm4.h>
using namespace std;
vector<uint8_t> randomData(size_t len){
    vector<uint8_t> res(len);
    rand_bytes(res.data(),len);
    return res;
}
vector<uint8_t> sm2_encrypt2(const SM2_KEY* key,const vector<uint8_t>& in){
    size_t len=in.size()+SM2_MAX_CIPHERTEXT_SIZE;
    vector<uint8_t> res(len);
    sm2_encrypt(key,in.data(),in.size(),res.data(),&len);
    res.resize(len);
    return res;
}
vector<uint8_t> sm2_decrypt2(const SM2_KEY* key,const vector<uint8_t>& in){
    size_t len=in.size();
    vector<uint8_t> res(len);
    sm2_decrypt(key,in.data(),in.size(),res.data(),&len);
    res.resize(len);
    return res;
}
vector<uint8_t> sm4_encrypt2(
    const vector<uint8_t>& key,
    const vector<uint8_t>& iv,
    const vector<uint8_t>& aad,
    const vector<uint8_t>& plaintext,
    vector<uint8_t>& tag) {
    SM4_KEY sm4_key;
    sm4_set_encrypt_key(&sm4_key, key.data());
    vector<uint8_t> ciphertext(plaintext.size());
    tag.resize(SM4_GCM_MAX_TAG_SIZE);
    sm4_gcm_encrypt(
        &sm4_key,iv.data(),iv.size(),
        aad.data(),aad.size(),
        plaintext.data(),plaintext.size(),
        ciphertext.data(),
        tag.size(),tag.data()
    );
    return ciphertext;
}
vector<uint8_t> sm4_decrypt2(
    const vector<uint8_t>& key,
    const vector<uint8_t>& iv,
    const vector<uint8_t>& aad,
    const vector<uint8_t>& ciphertext,
    const vector<uint8_t>& tag) {
    SM4_KEY sm4_key;
    sm4_set_encrypt_key(&sm4_key, key.data());
    vector<uint8_t> plaintext(ciphertext.size());
    if(sm4_gcm_decrypt(
        &sm4_key,iv.data(),iv.size(),
        aad.data(),aad.size(),
        ciphertext.data(),ciphertext.size(),
        tag.data(),tag.size(),
        plaintext.data())!=1){
        throw runtime_error("GCM标签验证失败");
    }
    return plaintext;
}
vector<uint8_t> genSm4Key(string pass){
    vector<uint8_t> res,salt=randomData(32);
    pbkdf2_hmac_sm3_genkey(pass.c_str(),pass.size(),salt.data(),32,10000,16,res.data());
    return res;
}
vector<uint8_t> sm9_encrypt2(const SM9_ENC_MASTER_KEY* mpk,const vector<uint8_t>& in,const string& id){
    size_t len=in.size()+SM9_MAX_CIPHERTEXT_SIZE;
    vector<uint8_t> res(len);
    sm9_encrypt(mpk,id.c_str(),id.size(),in.data(),in.size(),res.data(),&len);
    res.resize(len);
    return res;
}
vector<uint8_t> sm9_decrypt2(const SM9_ENC_KEY* key,const vector<uint8_t>& in,const string& id){
    size_t len=in.size();
    vector<uint8_t> res(len);
    sm9_decrypt(key,id.data(),id.size(),in.data(),in.size(),res.data(),&len);
    res.resize(len);
    return res;
}
void printHex(const vector<uint8_t>& in){
    cout<<hex<<setfill('0');
    for(auto& i:in) cout<<setw(2)<<(short)i;
    cout<<dec<<setfill(' ')<<endl;
}
int main() {
    cout<<"sm2 start"<<endl;
    SM2_KEY sm2_key;
    sm2_key_generate(&sm2_key);
    cout<<"public key:"<<endl;
    sm2_public_key_info_to_pem(&sm2_key,stdout);
    cout<<"private key:"<<endl;
    sm2_private_key_info_encrypt_to_pem(&sm2_key,"Password is 114514919810,***",stdout);
    cout<<"input:";
    string message;
    getline(cin,message);
    vector<uint8_t> enc=sm2_encrypt2(&sm2_key,vector<uint8_t>(message.begin(),message.end()));
    cout<<"Enc length:"<<enc.size()<<endl;
    cout<<"Enc:";
    printHex(enc);
    vector<uint8_t> dec=sm2_decrypt2(&sm2_key,enc);
    string res=string(dec.begin(),dec.end());
    cout<<"Dec length:"<<res.size()<<endl;
    cout<<"Dec:"<<res<<endl;

    cout<<"\nsm9 start"<<endl;
    string id="Id is 114514919810,***";
    cout<<"id:"<<id<<endl;
    SM9_ENC_MASTER_KEY sm9_msk;
    SM9_ENC_KEY sm9_key;
    sm9_enc_master_key_generate(&sm9_msk);
    sm9_enc_master_key_extract_key(&sm9_msk,id.c_str(),id.size(),&sm9_key);
    cout<<"master key(password=\"Password is 114514919810,***\"):"<<endl;
    sm9_enc_master_key_info_encrypt_to_pem(&sm9_msk,"Password is 114514919810,***",stdout);
    cout<<"key(password=\"Password is 114514919810,***\"):"<<endl;
    sm9_enc_key_info_encrypt_to_pem(&sm9_key,"Password is 114514919810,***",stdout);
    cout<<"input:";
    getline(cin,message);
    enc=sm9_encrypt2(&sm9_msk,vector<uint8_t>(message.begin(),message.end()),id);
    cout<<"Enc length:"<<enc.size()<<endl;
    cout<<"Enc:";
    printHex(enc);
    dec=sm9_decrypt2(&sm9_key,enc,id);
    res=string(dec.begin(),dec.end());
    cout<<"Dec length:"<<res.size()<<endl;
    cout<<"Dec:"<<res<<endl;

    cout<<"\nsm4 start"<<endl;
    vector<uint8_t> key=randomData(16);
    vector<uint8_t> iv=randomData(12);
    vector<uint8_t> aad=randomData(12);
    cout<<"key:";
    printHex(key);
    cout<<"iv:";
    printHex(iv);
    cout<<"aad:";
    printHex(aad);
    cout<<"input:";
    getline(cin,message);
    vector<uint8_t> tag;
    enc=sm4_encrypt2(key,iv,aad,vector<uint8_t>(message.begin(),message.end()),tag);
    cout<<"Enc length:"<<enc.size()<<endl;
    cout<<"Enc:";
    printHex(enc);
    dec=sm4_decrypt2(key,iv,aad,enc,tag);
    res=string(dec.begin(),dec.end());
    cout<<"Dec length:"<<dec.size()<<endl;
    cout<<"Dec: "<<res<<endl;
    return 0;
}